package cmd

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"syscall"

	"github.com/fatih/color"
	"github.com/linchao/zkp_auth/internal/client"
	"github.com/linchao/zkp_auth/models"
	"github.com/spf13/cobra"
)

var (
	user      string
	password  string
	out       string
	sessionId string
	deviceId  string
	dbPath    string
)

func SetupFlags() {
	RootCmd.PersistentFlags().StringVarP(&user, "user", "u", "", "User")
	RootCmd.PersistentFlags().StringVarP(&password, "password", "p", "", "Password")
	RootCmd.PersistentFlags().StringVarP(&out, "out", "o", "", "Output file for token")

	// 添加session_id验证相关的标志
	verifyCmd.Flags().StringVarP(&sessionId, "sid", "s", "", "Session ID to verify")
	verifyCmd.Flags().StringVarP(&deviceId, "did", "d", "", "Device ID for verification")
	verifyCmd.Flags().StringVarP(&dbPath, "db", "b", "./data/auth.db", "Path to SQLite database file")

	RootCmd.AddCommand(registerCmd)
	RootCmd.AddCommand(loginCmd)
	RootCmd.AddCommand(verifyCmd)
}

var RootCmd = &cobra.Command{
	Use:   "zkp_auth",
	Short: "A CLI for ZKP Authentication",
	Run: func(cmd *cobra.Command, args []string) {

		color.Yellow("************************ Welcome to Chaum-Pedersen ZKP based Authentication CLI *****************")
		color.Yellow("Please use 'register' or 'login' subcommands.")
		color.Yellow("To register a new user run: `go run main.go register -u <username> -p <password>`")
		color.Yellow("To login a registered user run: `go run main.go login -u <username> -p <password>`")
		color.Yellow("To save token to file, add -o flag: `go run main.go login -u <username> -p <password> -o <output_file>`")
		color.Yellow("To verify a session ID run: `go run main.go verify -d <device_id> -s <session_id>`")
		color.Yellow("To exit this terminal press CTRL+C")

		// Setup a signal handler to capture interrupt and termination signals
		done := make(chan os.Signal, 1)
		signal.Notify(done, os.Interrupt, syscall.SIGTERM)

		// when done exit the program gracefully
		<-done

	},
}

var registerCmd = &cobra.Command{
	Use:   "register",
	Short: "Register a new user",
	Run: func(cmd *cobra.Command, args []string) {
		grpcClient, err := client.SetupGRPCClient()
		if err != nil {
			log.Fatalf("error setting up grpc client %s", err.Error())
		}
		regRes, err := client.Register(*grpcClient, user, password)
		if err != nil {
			return
		}

		resJSON, err := json.Marshal(regRes)
		if err != nil {
			log.Fatal("error:", err)
		}

		// 输出结果
		color.Green(string(resJSON))

		// 如果指定了输出文件，则将结果写入文件
		if out != "" {
			if err := ioutil.WriteFile(out, resJSON, 0644); err != nil {
				log.Fatalf("保存到文件失败: %v", err)
			}
			color.Green(fmt.Sprintf("结果已保存到文件: %s", out))
		}
	},
}

var loginCmd = &cobra.Command{
	Use:   "login",
	Short: "Log in with a registered user",
	Run: func(cmd *cobra.Command, args []string) {
		grpcClient, err := client.SetupGRPCClient()
		if err != nil {
			log.Fatalf("error setting up grpc client %s", err.Error())
		}
		loginRes, err := client.LogIn(*grpcClient, user, password)
		if err != nil {
			return
		}

		resJSON, err := json.Marshal(loginRes)
		if err != nil {
			log.Fatal("error:", err)
		}

		// 输出结果
		color.Green(string(resJSON))

		// 如果指定了输出文件，则将结果写入文件
		if out != "" {
			var resultMap map[string]interface{}
			filename := out
			if err := json.Unmarshal(resJSON, &resultMap); err == nil {
				if sid, ok := resultMap["session_id"].(string); ok && sid != "" {
					// 如果out以/结尾，拼接session_id.json
					if out[len(out)-1] == '/' || out[len(out)-1] == '\\' {
						// 确保目录存在
						dir := out
						if dir == "" {
							dir = "./token/"
						}
						if err := os.MkdirAll(dir, 0755); err != nil {
							log.Fatalf("创建目录失败: %v", err)
						}
						filename = dir + sid + ".json"
					}
				}
			}
			if err := ioutil.WriteFile(filename, resJSON, 0644); err != nil {
				log.Fatalf("保存到文件失败: %v", err)
			}
			color.Green(fmt.Sprintf("会话令牌已保存到文件: %s", filename))
		}
	},
}

var verifyCmd = &cobra.Command{
	Use:   "verify",
	Short: "Verify a session ID for a device",
	Run: func(cmd *cobra.Command, args []string) {
		// 验证必须的参数
		if deviceId == "" {
			log.Fatalf("必须提供设备ID (-d 或 --did)")
		}
		if sessionId == "" {
			log.Fatalf("必须提供会话ID (-s 或 --sid)")
		}

		// 初始化数据库连接
		if err := models.InitDB(dbPath); err != nil {
			log.Fatalf("初始化数据库失败: %v", err)
		}
		defer models.CloseDB()

		// 查询设备
		device, err := models.GetDevice(context.Background(), deviceId)
		if err != nil {
			log.Fatalf("查询设备失败: %v", err)
		}

		// 检查设备是否存在
		if device == nil {
			color.Red("验证失败: 设备不存在")
			os.Exit(1)
		}

		// 检查session_id是否匹配
		if device.SessionID != sessionId {
			color.Red("验证失败: 会话ID不匹配")
			os.Exit(1)
		}

		// 验证成功
		color.Green("验证成功: 会话ID有效")

		// 构造响应对象
		type VerifyResponse struct {
			Success bool   `json:"success"`
			Message string `json:"message"`
		}

		resp := VerifyResponse{
			Success: true,
			Message: "会话ID有效",
		}

		respJSON, err := json.Marshal(resp)
		if err != nil {
			log.Fatal("错误:", err)
		}

		// 如果有指定输出文件，则写入文件
		if out != "" {
			if err := ioutil.WriteFile(out, respJSON, 0644); err != nil {
				log.Fatalf("保存到文件失败: %v", err)
			}
			color.Green(fmt.Sprintf("验证结果已保存到文件: %s", out))
		}
	},
}
